#include "oqs_prov.h"
#include "test_common.h"
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <openssl/core_dispatch.h>
#include <openssl/core_names.h>
#include <openssl/provider.h>
#include <openssl/ssl.h>
#include <openssl/pem.h>
#include <openssl/x509.h>
#include <openssl/x509v3.h>
#include <openssl/trace.h>
#include <openssl/params.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/param_build.h>
#include <openssl/bio.h> 
#include <openssl/asn1.h>
#include <openssl/ec.h>
#include <openssl/rsa.h>
#include <errno.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <time.h> 
const char *kHybridSignatureAlgorithms[] = {
    "p256_dilithium2", "rsa3072_dilithium2", "p384_dilithium3", "p521_dilithium5",
    "p256_mldsa44", "rsa3072_mldsa44", "p384_mldsa65", "p521_mldsa87",
    "p256_falcon512", "rsa3072_falcon512", "p256_falconpadded512", "rsa3072_falconpadded512",
    "p521_falcon1024", "p521_falconpadded1024", "p256_sphincssha2128fsimple", "rsa3072_sphincssha2128fsimple",
    "p256_sphincssha2128ssimple", "rsa3072_sphincssha2128ssimple", "p384_sphincssha2192fsimple",
    "p256_sphincsshake128fsimple", "rsa3072_sphincsshake128fsimple", NULL
};

struct SigPair {
    unsigned char *sigma1;
    size_t sigma1_len;
    unsigned char *sigma2;
    size_t sigma2_len;
} signature;

         
static int nid = NID_undef; 

static OSSL_LIB_CTX *init_openssl(void) {
    OSSL_LIB_CTX *ctx;

    if (!(ctx = OSSL_LIB_CTX_new())) {
        fprintf(stderr, cRED "failed to initialize a new `OSSL_LIB_CTX` " cNORM "\n");
    }
    return ctx;
}

void initialize_openssl_globals() {     
    signature.sigma1 = NULL;
    signature.sigma2 = NULL;
    signature.sigma1_len = 0;
    signature.sigma2_len = 0;

    nid = OBJ_create("1.2.3.4.1.2.3.4", "sigma1", "Sigma1");  
          if (nid == NID_undef) {           
         fprintf(stderr, "Initialize failed.\n"); 
   } 
}

static int is_signature_algorithm_hybrid(const char **list, const char *s) {
    for (; *list != NULL && strcmp(*list, s) != 0; ++list);
    return *list != NULL;
}

char *get_oqs_fromtls(const char *tlsname){
    char *oqsname;
    char *del = "_";
    char *position = strchr(tlsname, *del);
    
    if(position == NULL)
    return NULL;
    
    oqsname = position+1;
    
    return oqsname;
}

char *get_cls_fromtls(const char *tlsname) {
    char *clsname;
    char *del = "_";
    char *position = strchr(tlsname, *del);

    if (position == NULL) {
           return NULL;
    }

    size_t len = position - tlsname;

    clsname = (char *)malloc(len + 1);  // +1 为了 '\0'

    if (clsname == NULL) {
        return NULL;
    }
    
        strncpy(clsname, tlsname, len);
    clsname[len] = '\0';

    return clsname;
}

static EVP_PKEY *generate_classic_evp(const char *classic_names){
         
     EVP_PKEY *pkey = NULL;
            if(!strncmp(classic_names,"rsa", 3)){
        //        printf("1\n");
        pkey = EVP_RSA_gen(3072);
                    return pkey;
     }
           if(!strncmp(classic_names, "p256",4)){
   //                        printf("2\n");
         pkey =  EVP_EC_gen("P-256");
                       return pkey;
          }
            if(!strncmp(classic_names, "p384",4)){
   //                         printf("3\n");
         pkey =  EVP_EC_gen("P-384");
                       return pkey;
          }        
            if(!strncmp(classic_names, "p521",4)){
  //                          printf("4\n");
         pkey =  EVP_EC_gen("P-521");
                        return pkey;
          }          

    EVP_PKEY_free(pkey);
          return NULL;
}

static EVP_PKEY *generate_pq_evp(OSSL_LIB_CTX *libctx, const char *oqsname){
    EVP_PKEY_CTX *ctx = NULL;
    EVP_PKEY *pqkey = NULL;
 
   if (!(ctx = EVP_PKEY_CTX_new_from_name(libctx, oqsname, NULL)) ||!EVP_PKEY_keygen_init(ctx) ||!(EVP_PKEY_generate(ctx, &pqkey))) {
      fprintf(stderr, "X509 generate error.\n");
      EVP_PKEY_CTX_free(ctx);
      EVP_PKEY_free(pqkey);
               return NULL;
    }

    EVP_PKEY_CTX_free(ctx);
        return pqkey;    
}

X509_REQ *generate_cert_req(EVP_PKEY *pkey){
    X509_REQ *cert_req = X509_REQ_new();
    X509_NAME *name = NULL;
          if(!cert_req){
          fprintf(stderr, "X509_REQ_new error.\n");
          X509_REQ_free(cert_req);
          X509_NAME_free(name);
                         return NULL;
    }
    
    if(!X509_REQ_set_pubkey(cert_req, pkey)||!(name = X509_REQ_get_subject_name(cert_req))||!X509_NAME_add_entry_by_txt(name, "C",  MBSTRING_ASC, (unsigned char *)"CN", -1, -1, 0)||!X509_NAME_add_entry_by_txt(name, "O",  MBSTRING_ASC, (unsigned char *)"icatest.org", -1, -1, 0)||!X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char *)"ICA", -1, -1, 0)||!    X509_REQ_sign(cert_req, pkey, EVP_sha256())){
          fprintf(stderr, "X509_REQ_gen error.\n");
          X509_REQ_free(cert_req);
          X509_NAME_free(name);
                          return NULL;
    }

        return cert_req;    
} 

static  int get_signature_data(X509_REQ *cert_req, EVP_PKEY *pkey){
           size_t ret = 0;
    X509 *cert = X509_new();
    X509_NAME *name = X509_NAME_new();
    
   if(!X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char *)"CA", -1, -1, 0)||
  !X509_NAME_add_entry_by_txt(name, "O",  MBSTRING_ASC, (unsigned char *)"catest.org", -1, -1, 0)||!X509_NAME_add_entry_by_txt(name, "C",  MBSTRING_ASC, (unsigned char *)"CN", -1, -1, 0)){
        fprintf(stderr, "X509 NAME gen error.\n");
                     goto end;
    }    
    ASN1_INTEGER_set(X509_get_serialNumber(cert),1);
    X509_gmtime_adj(X509_getm_notBefore(cert), 0);
    X509_gmtime_adj(X509_getm_notAfter(cert), 31536000L);
    
    X509_set_subject_name(cert, X509_REQ_get_subject_name(cert_req));
    X509_set_issuer_name(cert, name);
    X509_set_pubkey(cert, X509_REQ_get_pubkey(cert_req)); 
          if(!X509_sign(cert, pkey,EVP_sha256())){
        fprintf(stderr, "X509_gen error.\n");        
                     goto end;
    }
          //fprintf(stderr, "X509_gen success.\n");
 
          const ASN1_BIT_STRING *sig;
          const X509_ALGOR *sigalg = NULL;  
    X509_get0_signature(&sig, &sigalg, cert);
          size_t siglen = 0;
    siglen = sig->length;
    signature.sigma1_len = siglen;
    signature.sigma1 = (unsigned char *)malloc(siglen);
         if (signature.sigma1== NULL) {
        fprintf(stderr, "Memory allocation failed\n");
                      goto end;
    } 
    memcpy(signature.sigma1,sig->data, siglen); 
    ret = 1;
end:
    X509_free(cert);
          return ret;    
}

X509 *generate_final_cert(X509_REQ *cert_req,EVP_PKEY *pkey){
    X509 *cert = X509_new();
    ASN1_INTEGER_set(X509_get_serialNumber(cert),1);
    X509_gmtime_adj(X509_getm_notBefore(cert), 0);
    X509_gmtime_adj(X509_getm_notAfter(cert), 31536000L);
    X509_set_subject_name(cert, X509_REQ_get_subject_name(cert_req));
    X509_set_pubkey(cert, X509_REQ_get_pubkey(cert_req)); 
    X509_NAME *name = X509_NAME_new(); 
    
       if(!X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char *)"CA", -1, -1, 0)||
  !X509_NAME_add_entry_by_txt(name, "O",  MBSTRING_ASC, (unsigned char *)"catest.org", -1, -1, 0)||!X509_NAME_add_entry_by_txt(name, "C",  MBSTRING_ASC, (unsigned char *)"CN", -1, -1, 0)){
        fprintf(stderr, "X509 NAME gen error.\n");
                     goto end;
    }
    X509_set_issuer_name(cert, name);
          
    X509_EXTENSION *sigma1 = NULL;
    ASN1_OCTET_STRING *data = ASN1_OCTET_STRING_new();
    
   // X509V3_EXT_add_alias(nid, NID_sigma_ext);
          if(!ASN1_OCTET_STRING_set(data, signature.sigma1, signature.sigma1_len)){
        fprintf(stderr, "Failed to create ASN1_OCTET_STRING for extension value.\n");
                      goto end;
    }
    
           if(!(sigma1 = X509_EXTENSION_create_by_NID(&sigma1, NID_subject_alt_name  ,1 ,data))){
        fprintf(stderr, "X509_EXTENSION create failed. \n");
                      goto end;
    }

    X509_add_ext(cert, sigma1, 0);
          if(cert == NULL){
      }
          if(!X509_sign(cert, pkey,EVP_sha256())){
        fprintf(stderr, "X509 final gen error.\n");        
                     goto end;
    }
    
//    fprintf(stderr,"Sigma2 gen success.\n"); 
    
         return cert;
end:
        X509_free(cert);
        X509_EXTENSION_free(sigma1);
   ASN1_OCTET_STRING_free(data);
        return NULL;
}

static int write_privkey_to_file(const char *filename, EVP_PKEY *pkey) {  
    FILE *file = fopen(filename, "wb");  
    if (file == NULL) {  
        fprintf(stderr, "Failed to open file for writing");  
    }  
  
    if (!PEM_write_PrivateKey(file, pkey, NULL, NULL, 0, NULL, NULL)) {  
        fclose(file);  
        fprintf(stderr,"Failed to write EVP_PKEY to file");  
    }  
  
    fclose(file);  
          return 1;  
}

static int write_x509_to_file(const char *filename, X509 *cert) {  
    FILE *file = fopen(filename, "wb");  
         if (file == NULL) {  
    fprintf(stderr, "Failed to open file for writing\n"); 
       return 0; 
    }    
    if (!PEM_write_X509(file, cert)) {  
        fprintf(stderr, "Failed to write X509 to file\n"); 
        fclose(file);  
               return 0; 
    }  
  
  fclose(file);  
    return 1;  
}  

int test_algorithm(OSSL_LIB_CTX *libctx,const char *signames,char * key_addr){

    EVP_PKEY *class_pkey = NULL;
    EVP_PKEY *pq_pkey = NULL;
          int ret = 1;
          char *oqsname, *clasname;
          char classkeypath[300];
          char oqskeypath[300];
          char certpath[300];
    signature.sigma1 = NULL;
    signature.sigma2 = NULL;
    signature.sigma1_len = 0;
    signature.sigma2_len = 0;

    oqsname = get_oqs_fromtls((char *)signames);
    clasname = get_cls_fromtls((char *)signames);
    
    //
       
    class_pkey  = generate_classic_evp(clasname);
    pq_pkey = generate_pq_evp(libctx, oqsname);
   
    X509_REQ *cert_req = X509_REQ_new();
    clock_t start_time = clock();  
          unsigned char* sigma1;
/*    cert_req = generate_cert_req(class_pkey); 
  //      cert_req = generate_cert_req(pq_pkey);
          if(get_signature_data(cert_req, class_pkey)){
//       fprintf(stderr,"Sigma1 gen success.\n");
    }else{
       fprintf(stderr,"Sigma1 gen failed.\n");  
                   goto end;  
    }*/
    cert_req = generate_cert_req(pq_pkey);
    X509 *final_cert = X509_new();
           if(get_signature_data(cert_req, class_pkey)){
//       fprintf(stderr,"Sigma1 gen success.\n");
    }else{
       fprintf(stderr,"Sigma1 gen failed.\n");  
                   goto end;  
    }    

    final_cert = generate_final_cert(cert_req, pq_pkey);
          if(final_cert == NULL){
       X509_free(final_cert);
                  goto end;
    }
    clock_t end_time = clock();     
    clock_t total_time = end_time - start_time;  
          double microseconds = ((double)total_time) * (1000000.0 / CLOCKS_PER_SEC);          
    fprintf(stderr, "程序运行时间: %.0f 微秒\n", microseconds);
    
          const char *sep = "/";
    
    sprintf(classkeypath, "%s%s%s%s", key_addr, sep, signames, "_cls.key");
    sprintf(oqskeypath, "%s%s%s%s", key_addr, sep, signames, "_oqs.key");
    sprintf(certpath, "%s%s%s%s", key_addr, sep, signames, ".crt");  
    

    write_privkey_to_file(classkeypath, class_pkey);
    write_privkey_to_file(oqskeypath, pq_pkey);
    write_x509_to_file(certpath, final_cert);   

    X509_free(final_cert);   
    ret = 0;
end:
    X509_REQ_free(cert_req);
    EVP_PKEY_free(class_pkey);
    EVP_PKEY_free(pq_pkey);
         return ret;
}


int main(int argc, char *argv[]) {
  OSSL_LIB_CTX *libctx;
  OSSL_PROVIDER *oqs_provider;
    char key_addr[100] ;
    const char *modulename;
    const char *configfile;
    const OSSL_ALGORITHM *sigalgs;
    int query_nocache;  
    int ret = EXIT_FAILURE;
          
    if (!(libctx = init_openssl())) {
                   goto end;
  }

  T(argc == 3);
  modulename = argv[1];
  configfile = argv[2];
    
    if(!(oqs_provider = OSSL_PROVIDER_load(libctx, modulename))){
        fputs(cRED "  `oqs_provider` is NULL " cNORM "\n", stderr);
                      goto end;    
    }
   initialize_openssl_globals();   
   sigalgs = OSSL_PROVIDER_query_operation(oqs_provider, OSSL_OP_SIGNATURE,
                                            &query_nocache);
    
    fprintf(stdout, "Enter the path to the files:");
        if (fgets(key_addr, sizeof(key_addr), stdin) == NULL) {
        fprintf(stderr, cRED " Failed to read input." cNORM "\n");
                  goto unload_def_provider;
    }
    key_addr[strcspn(key_addr, "\n")] = 0;
    //printf("%s\n", key_addr);  
    
    if (sigalgs) {
          for (; sigalgs->algorithm_names != NULL; sigalgs++) {
                  if (!is_signature_algorithm_hybrid(kHybridSignatureAlgorithms, sigalgs->algorithm_names)) {
                                continue;
          }
                            if (test_algorithm(libctx,sigalgs->algorithm_names, key_addr)) {
           fprintf(stderr, cRED " Failed for %s " cNORM "\n", sigalgs->algorithm_names);
        } else {
           fprintf(stderr, cGREEN "%s succeeded" cNORM "\n", sigalgs->algorithm_names);       
        }
    }
     }
    ret = 0;
unload_def_provider:
    OSSL_PROVIDER_unload(oqs_provider);   
end:
    OSSL_LIB_CTX_free(libctx);
    return ret;
  
}
